from langfuse import Langfuse
import requests
import os
import json
import re
import time

class LangfusePromptManager:
    def __init__(self, public_key=None, secret_key=None, host=None, dify_api_key=None, dify_api_url=None):
        """
        初始化 Langfuse 客户端和 Dify 配置
        """
        self.langfuse = Langfuse(
            public_key=public_key or "pk-lf-ea07ebae-8acb-4e55-8c74-10615512b4c3",
            secret_key=secret_key or "sk-lf-4432bb59-ef86-47ed-a66e-ea916cbd8368",
            host=host or "https://langfuse-dev.chinabeego.com",
            sample_rate=1.0,
            environment="default"
        )
        
        # Dify 配置
        self.dify_api_key = dify_api_key or "app-0PDgGG370Da8hMNlzK4k7jHT"
        self.dify_api_url = dify_api_url or "https://dify-pre.chinabeego.com/v1"
    
    def customer_service_chat(self, user_message, mobile=None, session_id=None):
        """
        客户服务聊天 - 使用 Langfuse 提示词管理和 Dify API
        """
        trace = self.langfuse.trace(name="customer_service_chat")
        
        try:
            # 从 Langfuse 获取提示词，并传入变量
            prompt = self.langfuse.get_prompt(
                name="customer-service-assistant",
                version="1.0.0"
            ).compile({
                "mobile": mobile or (re.search(r'1[3-9]\\d{9}', user_message).group() if re.search(r'1[3-9]\\d{9}', user_message) else "")
            })
            
            # 创建生成记录
            generation = trace.generation(
                name="customer_service_response",
                model="dify-ai-agent",
                input={"user_message": user_message, "mobile": mobile},
                prompt=prompt
            )
            
            # 调用 Dify API
            response_content = self._call_dify_api(
                user_message=user_message,
                inputs={"mobile": mobile, "system_prompt": prompt},
                session_id=session_id
            )
            
            # 结束生成记录
            generation.end(
                output=response_content,
                usage={"prompt_tokens": len(user_message) // 4, "completion_tokens": len(response_content) // 4}
            )
            
            return response_content
            
        except Exception as e:
            # 记录错误
            trace.event(
                name="error",
                input={"error": str(e), "user_message": user_message}
            )
            raise e
    
    def _call_dify_api(self, user_message, inputs=None, session_id=None):
        """
        调用 Dify API 生成回复
        """
        if inputs is None:
            inputs = {}
        
        if session_id is None:
            session_id = f"session_{int(time.time())}"
        
        headers = {
            "Authorization": f"Bearer {self.dify_api_key}",
            "Content-Type": "application/json"
        }
        
        # 准备Dify应用的输入参数
        dify_inputs = {}
        if "mobile" in inputs:
            dify_inputs["mobile"] = inputs["mobile"]
        else:
            # 尝试从用户消息中提取手机号
            mobile_match = re.search(r'1[3-9]\d{9}', user_message)
            if mobile_match:
                dify_inputs["mobile"] = mobile_match.group()
            else:
                dify_inputs["mobile"] = "未提供手机号"
        
        # 添加其他输入参数
        for key, value in inputs.items():
            if key != "mobile":
                dify_inputs[key] = value
        
        payload = {
            "inputs": dify_inputs,
            "query": user_message,
            "response_mode": "streaming",
            "user": session_id
        }
        
        try:
            response = requests.post(
                f"{self.dify_api_url}/chat-messages", 
                json=payload, 
                headers=headers, 
                timeout=120, 
                stream=True
            )
            response.raise_for_status()
            
            # 处理流式响应
            assistant_reply = ""
            for line in response.iter_lines():
                if line:
                    line_str = line.decode('utf-8')
                    
                    # 处理 Server-Sent Events 格式
                    if line_str.startswith('data: '):
                        data_str = line_str[6:]  # 移除 'data: ' 前缀
                        if data_str.strip() == '[DONE]':
                            break
                        
                        try:
                            data = eval(data_str)  # 解析数据
                            if isinstance(data, dict) and 'answer' in data:
                                assistant_reply += data['answer']
                        except:
                            # 如果不是有效的JSON，跳过
                            continue
            
            # 如果没有获取到答案，使用默认值
            if not assistant_reply:
                assistant_reply = "抱歉，无法获取AI回复。"
            
            return assistant_reply
            
        except Exception as e:
            raise Exception(f"Dify API 调用失败: {str(e)}")
    
    def create_prompt(self, name, prompt_content, version="1.0.0", description=""):
        """
        创建新的提示词
        """
        try:
            prompt = self.langfuse.create_prompt(
                name=name,
                prompt=prompt_content,
                version=version,
                description=description
            )
            print(f"提示词 '{name}' 创建成功，版本: {version}")
            return prompt
        except Exception as e:
            print(f"创建提示词失败: {e}")
            return None
    
    def get_prompt(self, name, version=None):
        """
        获取提示词
        """
        try:
            if version:
                prompt = self.langfuse.get_prompt(name=name, version=version)
            else:
                prompt = self.langfuse.get_prompt(name=name)
            return prompt
        except Exception as e:
            print(f"获取提示词失败: {e}")
            return None
    
    def list_prompts(self):
        """
        列出所有提示词
        """
        try:
            prompts = self.langfuse.get_prompts()
            return prompts
        except Exception as e:
            print(f"获取提示词列表失败: {e}")
            return []
    
    def update_prompt(self, name, prompt_content, version, description=""):
        """
        更新提示词
        """
        try:
            prompt = self.langfuse.update_prompt(
                name=name,
                prompt=prompt_content,
                version=version,
                description=description
            )
            print(f"提示词 '{name}' 更新成功，版本: {version}")
            return prompt
        except Exception as e:
            print(f"更新提示词失败: {e}")
            return None
    
    def delete_prompt(self, name, version):
        """
        删除提示词
        """
        try:
            self.langfuse.delete_prompt(name=name, version=version)
            print(f"提示词 '{name}' 版本 {version} 删除成功")
            return True
        except Exception as e:
            print(f"删除提示词失败: {e}")
            return False
    
    def chat_with_prompt(self, prompt_name, user_message, variables=None, session_id=None):
        """
        使用指定提示词进行聊天（通过 Dify API）
        """
        trace = self.langfuse.trace(name=f"chat_with_{prompt_name}")
        
        try:
            # 获取提示词
            prompt_obj = self.langfuse.get_prompt(name=prompt_name)
            
            # 编译提示词（如果有变量）
            mobile_var = None
            if variables and isinstance(variables, dict):
                mobile_var = variables.get("mobile")
            if not mobile_var:
                match = re.search(r'1[3-9]\\d{9}', user_message)
                mobile_var = match.group() if match else ""
            compiled_prompt = prompt_obj.compile({"mobile": mobile_var})
            
            # 创建生成记录
            generation = trace.generation(
                name=f"{prompt_name}_response",
                model="dify-ai-agent",
                input={"user_message": user_message, "mobile": mobile_var},
                prompt=compiled_prompt
            )
            
            # 调用 Dify API
            response_content = self._call_dify_api(
                user_message=user_message,
                inputs={"mobile": mobile_var, "system_prompt": compiled_prompt},
                session_id=session_id
            )
            
            # 结束生成记录
            generation.end(
                output=response_content,
                usage={"prompt_tokens": len(user_message) // 4, "completion_tokens": len(response_content) // 4}
            )
            
            return response_content
            
        except Exception as e:
            trace.event(
                name="error",
                input={"error": str(e), "user_message": user_message}
            )
            raise e

# 使用示例
def main():
    # 初始化提示词管理器
    prompt_manager = LangfusePromptManager()
    
    # 示例: 使用提示词进行聊天（仅使用 mobile 作为变量）
    response = prompt_manager.customer_service_chat(
        user_message="我想退货，手机号是13800138000",
        mobile="13800138000",
        session_id="test_session_001"
    )
    print("客户服务回复:", response)
    
    # 列出所有提示词
    prompts = prompt_manager.list_prompts()
    print("所有提示词:", prompts)
    
    # 使用通用聊天方法
    response2 = prompt_manager.chat_with_prompt(
        prompt_name="customer-service-assistant",
        user_message="帮我查下会员积分，手机号13800138000",
        variables={"mobile": "13800138000"},
        session_id="test_session_002"
    )
    print("通用聊天回复:", response2)

if __name__ == "__main__":
    main()